import matplotlib.pyplot as plt
import numpy as np
import soundfile as sf

def calc_STE(signal, sampsPerFrame):
    nFrames       = int(len(signal) / sampsPerFrame)  
    E = np.zeros(nFrames)
    for frame in range(nFrames):
        startIdx = frame * sampsPerFrame
        stopIdx = startIdx + sampsPerFrame
        E[frame]=np.sum(signal[startIdx:stopIdx] ** 2)
    return E

def calc_ZCR(signal, sampsPerFrame):
    nFrames = int(len(signal) / sampsPerFrame)       
    ZCR = np.zeros(nFrames)
    for frame in range(nFrames):
        startIdx = frame * sampsPerFrame
        stopIdx = startIdx + sampsPerFrame
        signalframe = signal[startIdx:stopIdx]
        for k in range(1, len(signalframe)):
            ZCR[frame] += 0.5 * abs(np.sign(signalframe[k]) - np.sign(signalframe[k - 1]))
    return ZCR
 
 
#file_name = 'voiced_unvoiced_e.wav' 
file_name = 'wave_3.wav'   
e, fs_e = sf.read(file_name) 
print("fs_z", fs_e)

#file_name = 'voiced_unvoiced_z.wav'
file_name = 'wave_3.wav'
z, fs_z = sf.read(file_name)   
print("fs_z", fs_z)

###################################################
signal = e
sampsPerFrame = int(0.02 * fs_e) #20 ms

plt.figure(figsize=(12, 6))
plt.subplot(3,1,1)
plt.plot(signal)

plt.subplot(3,1,2)
plt.plot(calc_ZCR(signal, sampsPerFrame))
plt.title('(Short-term) zcr per block ($L_{\mathrm{Bl}}=' + 
          str(sampsPerFrame) + '$, which is ' + 
          str(sampsPerFrame/fs_e*1000) + 'ms @ $f_s=' + str(fs_e) +'$)')
plt.tight_layout() # automatically adjust padding to make space for titles

plt.subplot(3,1,3)
plt.plot(calc_STE(signal, sampsPerFrame))
plt.title('(Short-term) Energy per block ($L_{\mathrm{Bl}}=' + 
          str(sampsPerFrame) + '$, which is ' + 
          str(sampsPerFrame/fs_e*1000) + 'ms @ $f_s=' + str(fs_e) +'$)')
plt.tight_layout() # automatically adjust padding to make space for titles
#plt.show()

###################################################
signal = z
sampsPerFrame = int(0.02 * fs_z) #20 ms

plt.figure(figsize=(12, 6))
plt.subplot(3,1,1)
plt.plot(signal)

plt.subplot(3,1,2)
plt.plot(calc_ZCR(signal, sampsPerFrame))

plt.title('(Short-term) zcr per block ($L_{\mathrm{Bl}}=' + 
          str(sampsPerFrame) + '$, which is ' + 
          str(sampsPerFrame/fs_z*1000) + 'ms @ $f_s=' + str(fs_z) +'$)')
plt.tight_layout() # automatically adjust padding to make space for titles

plt.subplot(3,1,3)
plt.plot(calc_STE(signal, sampsPerFrame))
plt.title('(Short-term) Energy per block ($L_{\mathrm{Bl}}=' + 
          str(sampsPerFrame) + '$, which is ' + 
          str(sampsPerFrame/fs_z*1000) + 'ms @ $f_s=' + str(fs_z) +'$)')
plt.tight_layout() # automatically adjust padding to make space for titles
plt.show()